/**
* @file main.c
* @Name: the main function
* @author W. P. Yao, ywp@pku.edu.cn
* @version 0.0.1
* @date 2015-12-30
* @brief 使用HYPRE软件包求解稀疏矩阵
*/
#include "Head.h"
#include "Datatype.h"
#include "Read_Data.h"
#include "Solver.h"
#include "Result.h"
#include "Mem_Free.h"


/**
* @Name: err_handler 
*
* @Param: comm
* @Param: err_code
* @Brief: comment on the error message when debug code 
*/
void err_handler(MPI_Comm * comm, int * err_code, ...)
{
	int resultlen;
	char err_string[MPI_MAX_ERROR_STRING];
	MPI_Error_string(*err_code, err_string, &resultlen);
	printf("MPI exception caughted: error code %d\n", *err_code);
	printf("\terror: %s\n", err_string);
	MPI_Abort(*comm, *err_code);
}

/**
* @Name: main 
*
* @Param: argc
* @Param: argv[]
*
* @Returns: 
*/
int main(int argc, char * argv[])
{
    // solver and preconditioner
    // solver_id:  1 -- AMG
    //             2 -- CG
    //             3 -- CG & AMG
    //             4 -- GMRES & AMG
    //             5 -- CG & ParaSail
    Solver_id = 2;
    
    // matrix size
    // file_number: 1 -- 14081
    //              2 -- 220673
    //              3 -- 879617
    file_number = 1;

	MPI_Init(&argc, &argv);							///MPI init
	MPI_Errhandler_create(&err_handler, &errhdl);
	MPI_Errhandler_set(MPI_COMM_WORLD, errhdl);
	MPI_Comm_rank(MPI_COMM_WORLD, &myrank);		    ///rank ID
	MPI_Comm_size(MPI_COMM_WORLD, &sum_rank);       ///rank num
	
	starttime = MPI_Wtime();						///start time

    //===========================================
    // step 0: read data from file
    // int Data_length,      the length of matrix.dat
    // int Matrix_size,      the size of the matrix
    // int *Data_I,          the column of the matrix
    // int *Data_J,          the row of the matrix
    // double *Data_A_Value, the value of matrix A
    // double *Data_b_Value, the value of vector b
    //===========================================
    Read_Data();                
    //printf("\nstep 0: read data done\n");	
    //============================================
	// step 1: initialize matrix and vector
    //============================================
    rows_average = Matrix_size/sum_rank;
    nrows        = rows_average;

    if(myrank == sum_rank-1){
        nrows = nrows + Matrix_size%sum_rank;
        }
    
    ilower = myrank*rows_average;
    iupper = ilower+nrows-1;
    jlower = ilower;
    jupper = iupper;

    HYPRE_IJMatrixCreate(MPI_COMM_WORLD, ilower, iupper, jlower, jupper, &A);
    HYPRE_IJMatrixSetObjectType(A, HYPRE_PARCSR);
    HYPRE_IJMatrixInitialize(A);

    HYPRE_IJVectorCreate(MPI_COMM_WORLD, ilower, iupper, &b);
    HYPRE_IJVectorSetObjectType(b, HYPRE_PARCSR);
    HYPRE_IJVectorInitialize(b);

    HYPRE_IJVectorCreate(MPI_COMM_WORLD, ilower, iupper, &x);
    HYPRE_IJVectorSetObjectType(x, HYPRE_PARCSR);
    HYPRE_IJVectorInitialize(x);
    
    //printf("\nstep 1: init done\n");	
    //============================================
	// step 2: set parameters for A, b and x
    // parameters: 
    // ncols[k]: number of non-zero elements in k-th col
    // rows[k]: index of row in given rank
    // cols[k]: index of col in given rank
    // num: number of non-zero elements in given rank
    // A_value, B_value: value of non-zero elements
    // length_cols[k]: number of non-zero elements in k-th cols.
    // num_cols[k]: sum of non-zero element number in 1~k cols.
    //============================================
    length_cols = (int*)malloc(Matrix_size*sizeof(int));
    num_cols    = (int*)malloc((Matrix_size+1)*sizeof(int));

    k = 0; m1 = 0; m2 = 1; num_cols[0] = 0;
    
    for (i = 0; i <  Data_length-1; i++){  // ???
        k = k+1;
        if (Data_I[i] != Data_I[i+1]){
            length_cols[m1] = k;
            num_cols[m2] = num_cols[m2-1]+k;
            m1 += 1; m2 += 1;
            k = 0;
            }
        }
    length_cols[m1] = k+1;
    num_cols[m2] = num_cols[m2-1]+k+1;

    ncols = (int*)malloc(nrows*sizeof(int)); 
     rows = (int*)malloc(nrows*sizeof(int));
    for (i = 0; i < nrows; i++){
        k = myrank*rows_average+i;
        rows[i] = k;
        ncols[i] = length_cols[k];
        }
    
    num  = num_cols[iupper+1] - num_cols[ilower];
    cols = (int*)malloc(num*sizeof(int));
    A_value = (double*)malloc(num*sizeof(double));
    for (i = 0; i < num; i++){
        k = num_cols[ilower]+i;
        cols[i] = Data_J[k];
        A_value[i] = Data_A_Value[k];
        }

    b_value = (double*)malloc(nrows*sizeof(double));
    for (i = 0; i < nrows; i++){
        k = ilower + i;
        b_value[i] = Data_b_Value[k];
        }
    
    HYPRE_IJMatrixSetValues(A, nrows, ncols, rows, cols, A_value);
    HYPRE_IJVectorSetValues(b, nrows, rows, b_value);
    HYPRE_IJMatrixPrint(A,"IJ.out.A");

    //printf("\nstep 2: settings done\n");	

    //=======================================
    // step 3: make sure they are ready
    //======================================
    HYPRE_IJMatrixAssemble(A);
    HYPRE_IJVectorAssemble(b);
    HYPRE_IJVectorAssemble(x);


    //printf("\nstep 3: ready for use\n");	
    //=======================================
    // step 4: get object 
    //======================================
    HYPRE_ParCSRMatrix A0;
    HYPRE_IJMatrixGetObject(A, (void **)(&A0));
    HYPRE_ParVector b0;
    HYPRE_IJVectorGetObject(b, (void **)(&b0));
    HYPRE_ParVector x0;
    HYPRE_IJVectorGetObject(x, (void **)(&x0));

    //printf("\nstep 4: done\n");	

    //=======================================
    // step 5: Solve Ax = b 
    //======================================
    if(Solver_id == 1){
        AMG(A0, b0, x0);
        }
    else if(Solver_id == 2){
        CG(A0, b0, x0);
        }
    else if(Solver_id == 3){
        CG_AMG(A0, b0, x0);
        }
    else if(Solver_id == 4){
        GMRES_AMG(A0, b0, x0);
        }
    else if(Solver_id == 5){
        CG_ParaSails(A0, b0, x0);
        }
    

    //printf("\nstep 5: done\n");	
    //=======================================
    // step 6: output data 
    // HYPRE_IJMatrixSetPrintLevel(A,3);
    // HYPRE_IJMatrixPrint(A, "./result/A");
    // HYPRE_IJVectorPrint(b, "./result/b");
    // HYPRE_IJVectorPrint(x, "./result/x");
    //======================================
    Result(x);

    //printf("\nstep 6: done\n");	
    

	endtime  = MPI_Wtime();
	duration = endtime - starttime;				///total run time
	MPI_Finalize();									///end of MPI

	if(myrank == 0){
		printf("\nThe run time is %f\n", duration); ///print run time
	}
    
    Mem_Free();
	return 0;
}
